package pool

import "sync"

var (
	pool4K   sync.Pool
	pool8K   sync.Pool
	pool16K  sync.Pool
	pool32K  sync.Pool
	pool64K  sync.Pool
	pool128K sync.Pool
	pool256K sync.Pool
	pool512K sync.Pool
	pool1M   sync.Pool
	pool2M   sync.Pool
	pool4M   sync.Pool
)

func bytesMallocByPool(pool *sync.Pool, size int) []byte {
	if v := pool.Get(); v != nil {
		return v.([]byte)
	}

	return make([]byte, size)
}

func bytesMalloc(size int) []byte {
	if size <= s4K {
		return bytesMallocByPool(&pool4K, s4K)
	}

	if size <= s8K {
		return bytesMallocByPool(&pool8K, s8K)
	}

	if size <= s16K {
		return bytesMallocByPool(&pool16K, s16K)
	}

	if size <= s32K {
		return bytesMallocByPool(&pool32K, s32K)
	}

	if size <= s64K {
		return bytesMallocByPool(&pool64K, s64K)
	}

	if size <= s128K {
		return bytesMallocByPool(&pool128K, s128K)
	}

	if size <= s256K {
		return bytesMallocByPool(&pool256K, s256K)
	}

	if size <= s512K {
		return bytesMallocByPool(&pool512K, s512K)
	}

	if size <= s1M {
		return bytesMallocByPool(&pool1M, s1M)
	}

	if size <= s2M {
		return bytesMallocByPool(&pool2M, s2M)
	}

	if size <= s4M {
		return bytesMallocByPool(&pool4M, s4M)
	}

	return make([]byte, size)
}

func BytesMalloc(size int) []byte {
	return bytesMalloc(size)[:size]
}

func BytesFree(bytes []byte) {
	size := cap(bytes)
	bytes = bytes[:size]

	if size == s4K {
		pool4K.Put(bytes)
		return
	}

	if size == s8K {
		pool8K.Put(bytes)
		return
	}

	if size == s16K {
		pool16K.Put(bytes)
		return
	}

	if size == s32K {
		pool32K.Put(bytes)
		return
	}

	if size == s64K {
		pool64K.Put(bytes)
		return
	}

	if size == s128K {
		pool128K.Put(bytes)
		return
	}

	if size == s256K {
		pool256K.Put(bytes)
		return
	}

	if size == s512K {
		pool512K.Put(bytes)
		return
	}

	if size == s1M {
		pool1M.Put(bytes)
		return
	}

	if size == s2M {
		pool2M.Put(bytes)
		return
	}

	if size == s4M {
		pool4M.Put(bytes)
		return
	}
}
